import numpy as np
from scipy import spatial

import torch
from torch import nn
from torchmetrics import Accuracy, MetricCollection, Precision, Recall,CosineSimilarity
from torchmetrics.functional.pairwise import pairwise_cosine_similarity
from sklearn.metrics import confusion_matrix

def zsl_acc(X,W,y_true,sig):
    """
    X: 特征
    W: 参数
    y_true: 标签值
    sig: 属性
    """
    y_true_union = np.unique(y_true)
    XW = torch.matmul(X,W)
    # 计算余弦相似度
    dist = 1- spatial.distance.cdist(XW,sig,'cosine')# N*C （no. f classes)
    # 也是计算余弦相似度，这里已经是用 1- 余弦相似度的结果了
    dist2 = pairwise_cosine_similarity(XW,sig)# N*C （no. f classes)
    softmax = torch.nn.Softmax(dim=1)
    outputs = softmax(torch.tensor(dist))
    preds = torch.max(outputs, dim=1)[1]
    # preds = np.array([np.argmax(output) for output in dist])
    cm = confusion_matrix(y_true,preds)
    sum_dim1 = cm.sum(axis=1)[:, np.newaxis]
    cm = cm.astype('float') / sum_dim1
    sss = sum(cm.diagonal())
    acc = sum(cm.diagonal()) / sig.shape[0]
    return acc

def zsl_acc2(attr_preds,y_true,sig):
    """
    X: 特征
    W: 参数
    y_true: 标签值
    sig: 属性
    """
    y_true_union = np.unique(y_true)
    # 计算余弦相似度
    dist = 1- spatial.distance.cdist(attr_preds,sig,'cosine')# N*C （no. f classes)
    # 也是计算余弦相似度，这里已经是用 1- 余弦相似度的结果了
    # dist2 = pairwise_cosine_similarity(attr_preds,sig)# N*C （no. f classes)
    softmax = torch.nn.Softmax(dim=1)
    outputs = softmax(torch.tensor(dist))
    preds = torch.max(outputs, dim=1)[1]
    preds_union = np.unique(preds)
    # preds = np.array([np.argmax(output) for output in dist])
    preds2 = torch.index_select(torch.from_numpy(y_true_union), dim=0,
                               index=preds)
    cm = confusion_matrix(y_true,preds2)
    sum_dim1 = cm.sum(axis=1)[:, np.newaxis]
    cm = cm.astype('float') / sum_dim1
    sss = sum(cm.diagonal())
    acc = sum(cm.diagonal()) / sig.shape[0]
    return acc
def zsl_acc_gzsl(X, W, y_true, classes, sig): # Class Averaged Top-1 Accuarcy
    y_true_union = np.unique(y_true)
    class_scores = torch.matmul(torch.matmul(X, W), sig.T)  # N x Number of Classes
    softmax = torch.nn.Softmax(dim=1)
    outputs = softmax(torch.tensor(class_scores))
    y_pred = torch.max(outputs, dim=1)[1]+1
    y_pred_union = np.unique(y_pred)
    per_class_acc = torch.zeros(len(classes))

    for i in range(len(classes)):
        is_class = y_true == classes[i]
        per_class_acc[i] = ((y_pred[is_class] == y_true[is_class]).sum()) / is_class.sum()

    return per_class_acc.mean()

def zsl_acc_gzsl2(attr_preds, y_true,classes, sig): # Class Averaged Top-1 Accuarcy
    y_true_union = np.unique(y_true)
    # 创建 CosineSimilarity 层
    dist = 1 - spatial.distance.cdist(attr_preds, sig, 'cosine')  # N*C （no. f classes)
    softmax = torch.nn.Softmax(dim=1)
    outputs = softmax(torch.tensor(dist))
    # 预测值
    y_pred = torch.max(outputs, dim=1)[1]
    y_pred_union = np.unique(y_pred)
    per_class_acc = torch.zeros(len(classes))

    for i in range(len(classes)):
        is_class = y_true == classes[i]
        sum = is_class.sum()
        ss = ((y_pred[is_class] == y_true[is_class]).sum())
        per_class_acc[i] = ((y_pred[is_class] == y_true[is_class]).sum()) / is_class.sum()

    return per_class_acc.mean()
class Evaluator():
    def __init__(self, num_classes) -> None:
        self._device = torch.device("cpu")
        self.hasAttrMetric = False
        self.hasLabelMetric = False
        # collection of all validation metrics
        self.metric_collection_attr = MetricCollection({
            'attr_acc': Accuracy(task="multiclass", num_classes=num_classes),
            'prec': Precision(task="multiclass", num_classes=num_classes, average='macro'),
            'rec': Recall(task="multiclass", num_classes=num_classes, average='macro')
        }).to(self._device)

        self.metric_collection_label = MetricCollection({
            'attr_acc': Accuracy(task="multiclass", num_classes=num_classes),
            'prec': Precision(task="multiclass", num_classes=num_classes, average='macro'),
            'rec': Recall(task="multiclass", num_classes=num_classes, average='macro')
        }).to(self._device)

    def evaluate_predictions(self, attr_preds, label_preds, attrs, labels):
        if attr_preds is not None:
            self.hasAttrMetric = True
            self.evaluate_att_predictions(attr_preds, attrs, labels)
        if label_preds is not None:
            self.hasLabelMetric = True
            self.evaluate_label_predictions(label_preds, labels)

    def evaluate_label_predictions(self, predictions, truth):
        predictions = predictions.to(self._device)
        truth = truth.to(self._device)
        y_true_union = np.unique(truth)
        softmax = torch.nn.Softmax(dim=1)
        outputs = softmax(predictions)
        preds = torch.max(outputs, dim=1)[1]
        y_preds_union = np.unique(preds)
        batch_acc = self.metric_collection_label(preds, truth)
        return batch_acc

    def evaluate_att_predictions(self, attr_preds, attrs, labels):
        """
            评估指标
            Inputs
                predictions:模型评分
                attrs: 所有属性矩阵
                truth: 真实标签
            Returns:
                batch_materic 指标值
        """
        attr_preds = attr_preds.to(self._device)
        labels = labels.to(self._device)
        attrs = torch.tensor(attrs).to(self._device)
        # 创建 CosineSimilarity 层
        cosine_similarity = nn.CosineSimilarity(dim=1, eps=1e-6)
        similaritys = []
        # 计算余弦相似度
        for prediction in attr_preds:
            prediction = prediction.repeat(len(attrs), 1)
            similarity = cosine_similarity(prediction, attrs)
            similaritys.append(similarity)
        similaritys = torch.stack(similaritys, dim=0)
        softmax = torch.nn.Softmax(dim=1)
        outputs = softmax(similaritys)
        # 预测值
        preds = torch.max(outputs, dim=1)[1]
        batch_materic = self.metric_collection_attr(preds, labels)
        return batch_materic

    def compute(self):
        metric_attr = None
        metric_label = None
        if self.hasAttrMetric:
            metric_attr = self.metric_collection_attr.compute()
        if self.hasLabelMetric:
            metric_label = self.metric_collection_label.compute()
        return metric_attr,metric_label

    def reset(self):
        self.metric_collection_attr.reset()
        self.metric_collection_label.reset()
